# -*- coding: utf-8 -*-

# File Name： res_layer
# Description :
# Author : lirui
# create_date： 2022/6/4
# Change Activity:
from torch.nn import Sequential


class ResLayer(Sequential):
    """
    ResLayer to build ResNet style backbone.
    """

    def __init__(self, block, in_channels, out_channels, num_blocks, stride=1,
                 norm_cfg=None, downsample_first=True, expansion=1, dilation=1, **kwargs):
        if norm_cfg is None:
            norm_cfg = dict(type='BN')
        self.block = block
        downsample = None
        self.expansion = expansion
        layers = []
        if downsample_first:
            layers.append(block(in_channels=in_channels,
                                out_channels=out_channels,
                                kernel_size=3,
                                stride=stride,
                                expansion=expansion,
                                dilation=dilation,
                                norm_cfg=norm_cfg,
                                **kwargs))
            for _ in range(1, num_blocks):
                layers.append(block(in_channels=out_channels,
                                    out_channels=out_channels,
                                    kernel_size=3,
                                    stride=1,
                                    dilation=dilation,
                                    norm_cfg=norm_cfg,
                                    **kwargs))

        else:  # downsample_first=False is for HourglassModule
            for _ in range(num_blocks - 1):
                layers.append(block(in_channels=in_channels,
                                    out_channels=in_channels,
                                    kernel_size=3,
                                    stride=1,
                                    expansion=expansion,
                                    dilation=dilation,
                                    norm_cfg=norm_cfg,
                                    **kwargs))
            layers.append(block(in_channels=in_channels,
                                out_channels=out_channels,
                                stride=stride,
                                dilation=dilation,
                                downsample=downsample,
                                norm_cfg=norm_cfg,
                                **kwargs))
        super(ResLayer, self).__init__(*layers)
